/* -*- mode: C++; c-file-style: "stroustrup"; c-basic-offset: 4; indent-tabs-mode: nil; -*- */

// Compile like this: g++ -O2 -Wall matrix.cpp -o matrix

#include <iostream>
#include <cstdlib>
#include <cassert>
#include <cstring>

#include <limits.h>
#include <math.h>

/* For testing, you can use a
 * constant value if you wish.
 */
#define RANGE(S) (S)

// Precision for checking identity.
#define PRECISION 1e-10

using namespace std;

/* Random generation, but always the same for
 * a given size for fair comparison.
 */
void gen_mat(double *a, size_t dim)
{
    srand48(dim);

    for(size_t i = 0; i < dim; ++i)
    {
        for(size_t j = 0; j < dim; ++j)
        {
            double z = (drand48() - 0.5)*RANGE(dim);
            a[i*dim+j] = (z < 10*PRECISION && z > -10*PRECISION) ? 0.0 : z;
        }
    }
}

void mat_mult(const double *a, const double *b, double *c, size_t dim)
{
    assert(a != c && b != c);

    /* Re-arranged matrix multiplication - variant. */

    for(size_t i = 0; i < dim; ++i)
    {
        for(size_t j = 0; j < dim; ++j)
        {
            c[i*dim+j] = 0;
        }
        for(size_t k = 0; k < dim; ++k)
        {
            for(size_t j = 0; j < dim; ++j)
            {
                c[i*dim+j] += a[i*dim+k]*b[k*dim+j];
            }
        }
    }    
}

/* Check that a divider is not too close to 0. */
void check_divider(double x)
{
    if (fabs(x) < PRECISION)
    {
        cerr << "Matrix non-invertible or computationally unstable.\n";
        abort();
    }
}

/* Read a re-arranged matrix. */
#define A(I,J) a[index[I]*n+(J)]
#define B(I,J) b[index[I]*n+(J)]

/* a = invert(input), using b as temporary matrix
 * Do not mind the signature, this is old code.
 * You do not need to look at the implementation in detail.
 */
void inv_mat(const double *input, double *a, double *b, size_t n)
{
    size_t index[n];

    // Init a, b, and index.
    memcpy(a, input, n*n*sizeof(double));
    memset(b, 0, n*n*sizeof(double));
    for(size_t i = 0; i < n; ++i)
    {
        index[i] = i;   // Identity mapping.
        b[i*n+i] = 1.0; // Identity matrix.
    }

    // Gaussian elimination.
    for(size_t k = 0; k < n; ++k)
    {
        // Find pivot.
        double x = fabs(A(k,k));
        size_t j = k;
        for(size_t i = k+1; i < n; ++i)
        {
            double y = fabs(A(i,k));
            if (y > x)
            {
                j = i;
                x = y;
            }
        }

        // Swap j <-> k.
        size_t i = index[j];
        index[j] = index[k];
        index[k] = i;

        // Division step.
        x = A(k,k);
        A(k,k) = 1.0;
        check_divider(x);
        for(size_t j = k+1; j < n; ++j) A(k,j) /= x;
        for(size_t j = 0  ; j < n; ++j) B(k,j) /= x;

        // Elimination step.
        for(i = k+1; i < n; ++i)
        {
            for(j = k+1; j < n; ++j) A(i,j) -= A(i,k)*A(k,j);
            for(j = 0  ; j < n; ++j) B(i,j) -= A(i,k)*B(k,j);
            A(i,k) = 0.0;
        }
    }

    // Back-substitution.
    for(size_t k = n-1; k > 0; --k)
    {
        for(size_t i = k-1; (int)i >= 0; --i)
        {
            for(size_t j = 0; j < n; ++j) B(i,j) -= A(i,k)*B(k,j);
            //A(i,k) = 0.0; -- implicit and not used later.
        }
    }

    // Result.
    for(size_t i = 0; i < n; ++i)
    {
        for(size_t j = 0; j < n; ++j)
        {
            a[i*n+j] = B(i,j);
        }
    }
}

void check_identity(const double *a, size_t dim)
{
    double error = 0.0;

    for(size_t i = 0; i < dim; ++i)
    {
        for(size_t j = 0; j < dim; ++j)
        {
            double r = i == j ? a[i*dim+j] - 1.0 : a[i*dim+j];
            if (r < -PRECISION || r > PRECISION)
            {
                cerr << "Matrix is not identity.\n";
                abort();
            }
            error += fabs(r);
        }
    }

    /* Smaller error = better precision. */
    cout << "\033[K\rError=" << error << endl;
}

void test(size_t dim)
{
    double *a = new double[dim*dim];
    double *b = new double[dim*dim];
    double *c = new double[dim*dim];

    gen_mat(a, dim);         // a = random matrix
    inv_mat(a, b, c, dim);   // b = invert(a)
    gen_mat(c, dim);         // randomize c
    mat_mult(a, b, c, dim); // c = a*b
    check_identity(c, dim);

    delete [] c;
    delete [] b;
    delete [] a;
}


int main(int argc, char *argv[])
{
    if (argc < 2)
    {
        cerr << "Usage: " << argv[0] << " size\n";
        return 1;
    }
    else
    {
        int i = atoi(argv[1]);
        if (i < 1)
        {
            cerr << "Minimum dimension is 1.\n";
            return 1;
        }
        test((size_t) i);
        cout << "Test passed.\n";
        return 0;
    }
}